#include "arp.h"
#include <signal.h>

struct list_head cache_list_head;
extern char my_ip_addr[];
extern char my_mac_addr[];
extern int my_if_index;

void sig_handler(int signo)
{
        printf("\nSIGINT received...\nARP module is now exiting...\n");
        unlink(DOM_SOCK_FILE);
        exit(0);
}

int main(void)
{
    int packet_socket, areq_socket, maxfd;
    int ret = 0;
    fd_set readfds;
    struct sockaddr_un domaddr;
    struct sockaddr_ll pfaddr;
    
    /* Ignore SIGPIPE */
    signal(SIGPIPE, SIG_IGN);
    signal(SIGINT, sig_handler);
    
    INIT_LIST_HEAD(&cache_list_head);

    /* Build a set of  <IP address , HW address>  matching pairs for all eth0
     * interface IP addresses (including alias IP addresses, if any).
     *
     * Print the address pairs found. 
     */
    print_hw_addrs();

    /* ARP request messages arrive on the PF_PACKET socket, the module
     * processes them, and responds with ARP reply messages as appropriate.
     */
    if((packet_socket = socket(PF_PACKET, SOCK_RAW, htons(ETH_ARP_PROTO))) < 0) {
        fprintf(stderr, "ARP: Error in creating PF_PACKET socket - %s\n", strerror(errno));
        ret = packet_socket;
        goto out;
    }

    memset(&pfaddr, 0, sizeof(struct sockaddr_ll));
    pfaddr.sll_family = PF_PACKET;
    pfaddr.sll_protocol = htons(ETH_ARP_PROTO);
    pfaddr.sll_ifindex = my_if_index;
    pfaddr.sll_hatype = ARPHRD_ETHER;
    pfaddr.sll_pkttype = PACKET_OTHERHOST;
    pfaddr.sll_halen = ETH_ALEN;
    memset(pfaddr.sll_addr, 0xff, ETH_ALEN);
    pfaddr.sll_addr[6] = 0x00;
    pfaddr.sll_addr[7] = 0x00;
    
    if(bind(packet_socket, (struct sockaddr*)&pfaddr, sizeof(pfaddr)) < 0) {
        fprintf(stderr, "ARP: Error in binding PF_PACKET socket - %s\n", strerror(errno));
        ret = -1;
        goto out;
    }

    /* This socket will be used to communicate with the function areq that is
     * implemented in the Tour module
     */
    if((areq_socket = socket(AF_LOCAL, SOCK_STREAM, 0)) < 0) {
        fprintf(stderr, "ARP: Error in creating DOMAIN socket - %s\n", strerror(errno));
        ret = areq_socket;
        close(packet_socket);
        goto out;
    }

    unlink(DOM_SOCK_FILE);

    memset(&domaddr, 0, sizeof(domaddr));
    domaddr.sun_family = AF_LOCAL;
    strcpy(domaddr.sun_path, DOM_SOCK_FILE);

    if (bind(areq_socket, (struct sockaddr *)&domaddr, sizeof(domaddr)) < 0) {
        fprintf(stderr, "ARP: Error in binding the DOMAIN socket - %s\n", strerror(errno));
        ret = -1;
        goto out_file;
    }

    if(listen(areq_socket, 5) < 0) {
        fprintf(stderr, "ARP: Error in creating a listening DOMAIN socket - %s\n", strerror(errno));
        ret = -1;
        goto out_file;
    }

    
    maxfd = packet_socket > areq_socket ? packet_socket : areq_socket;

    for(;;) {
        FD_ZERO(&readfds);
        FD_SET(packet_socket, &readfds);
        FD_SET(areq_socket, &readfds);

        if(select(maxfd + 1, &readfds, NULL, NULL, NULL) < 0) {
                if(errno != EINTR) {
                        fprintf(stderr, "ARP: Error in select() - %s\n", strerror(errno));
                        ret = -1;
                        goto out_file;
                }
        }

        if(FD_ISSET(packet_socket, &readfds)) {
            handle_message_on_packet_socket(packet_socket, &pfaddr);
        }
        else if(FD_ISSET(areq_socket, &readfds)) {
            handle_message_on_domain_socket(areq_socket, packet_socket, &pfaddr);
        }
    }
out_file:
    unlink(DOM_SOCK_FILE);

    close(packet_socket);
    close(areq_socket);
out:
    return ret;
} 

void handle_message_on_packet_socket(int packet_sockfd, struct sockaddr_ll *pfaddr)
{
    unsigned char * recv_buf = malloc(ETH_FRAME_LEN);
    unsigned char * send_buf = malloc(ETH_FRAME_LEN);
    int len = 0;
    int if_index = 0;
    struct address_cache *node = malloc(sizeof(struct address_cache));
    unsigned short hatype = 0;
    struct arp_packet *arp_msg;
    struct ethhdr *eh_send = (struct ethhdr*)send_buf;
    struct ethhdr *eh_recv = (struct ethhdr*)recv_buf;
    struct address_cache *temp;

    if(recv_buf == NULL || send_buf == NULL || node == NULL) {
        fprintf(stderr, "ARP: Failed to allocate memory.\n");
        exit(-1);
    }

    len = recvfrom(packet_sockfd, recv_buf, ETH_FRAME_LEN, 0, NULL, NULL);
    if(len < 0) {
        printf("ARP: Error reading from RAW socket.\n");
        goto out;
    }
    /* Position the pointer to the start of the op field in the ARP request/reply packet */
    arp_msg = (struct arp_packet *)(recv_buf + 14); 

    if_index = my_if_index;
    hatype = ARPHRD_ETHER;

    if(arp_msg->id != ARP_ID) {
        fprintf(stderr, "ARP: Unknown ARP_ID in received packet; ignoring this packet\n");
        goto out;   /* The 2-byte identification field did not match; dropping this packet */
    }

    if(arp_msg->type == ARP_REQUEST) {
            if(!memcmp(arp_msg->target_ip, my_ip_addr, IP_SIZE)) {
                    // printf("\nDEBUG: I am the destination node.\n");
                    char src_ipaddress[INET_ADDRSTRLEN + 1] = {0};
                    char dest_ipaddress[INET_ADDRSTRLEN + 1] = {0};
                    char vm_name[128] = {0};
                    struct sockaddr_ll unicast_pfaddr;
                    
                    memset(&unicast_pfaddr, 0, sizeof(struct sockaddr_ll));
                    unicast_pfaddr.sll_family = PF_PACKET;
                    unicast_pfaddr.sll_protocol = htons(ETH_ARP_PROTO);
                    unicast_pfaddr.sll_ifindex = my_if_index;
                    unicast_pfaddr.sll_hatype = ARPHRD_ETHER;
                    unicast_pfaddr.sll_pkttype = PACKET_OTHERHOST;
                    unicast_pfaddr.sll_halen = ETH_ALEN;

                    memcpy(unicast_pfaddr.sll_addr, arp_msg->sender_mac, IF_HADDR);

                    unicast_pfaddr.sll_addr[6] = 0x00;
                    unicast_pfaddr.sll_addr[7] = 0x00;

                    gethostname(vm_name, 128);
                    printf("\nReceived ARP_REQUEST message...\n");
                    printf("\nARP_REQUEST message received at %s(%s)::\n", vm_name, inet_ntop(AF_INET, my_ip_addr, src_ipaddress, INET_ADDRSTRLEN));
                    printf("Ethernet frame header -- src: %02x:%02x:%02x:%02x:%02x:%02x --> dst: %02x:%02x:%02x:%02x:%02x:%02x\n", 
                           (unsigned char)eh_recv->h_source[0], (unsigned char)eh_recv->h_source[1], (unsigned char)eh_recv->h_source[2], (unsigned char)eh_recv->h_source[3], (unsigned char)eh_recv->h_source[4], 
                           (unsigned char)eh_recv->h_source[5],
                           (unsigned char)eh_recv->h_dest[0], (unsigned char)eh_recv->h_dest[1], (unsigned char)eh_recv->h_dest[2], (unsigned char)eh_recv->h_dest[3], (unsigned char)eh_recv->h_dest[4],
                           (unsigned char)eh_recv->h_dest[5]);

                    printf("Payload -- sender_ip: %s, sender_mac: %02x:%02x:%02x:%02x:%02x:%02x, target_ip: %s, target_mac: %02x:%02x:%02x:%02x:%02x:%02x\n\n", 
                           inet_ntop(AF_INET, arp_msg->sender_ip, src_ipaddress, INET_ADDRSTRLEN), 
                           (unsigned char)arp_msg->sender_mac[0], (unsigned char)arp_msg->sender_mac[1], (unsigned char)arp_msg->sender_mac[2], (unsigned char)arp_msg->sender_mac[3], (unsigned char)arp_msg->sender_mac[4], (unsigned char)arp_msg->sender_mac[5],
                           inet_ntop(AF_INET, arp_msg->target_ip, dest_ipaddress, INET_ADDRSTRLEN),
                           (unsigned char)arp_msg->target_mac[0], (unsigned char)arp_msg->target_mac[1], (unsigned char)arp_msg->target_mac[2], (unsigned char)arp_msg->target_mac[3], (unsigned char)arp_msg->target_mac[4], (unsigned char)arp_msg->target_mac[5]);

                    /* I am the destination */
                    // printf("\nDEBUG: I am the destination for this ARP_REQUEST\n");
                    struct arp_packet *arp_reply;

                    arp_reply = (struct arp_packet *)(send_buf + 14);

                    /* Populate the ARP reply contents first */
                    arp_reply->type = ARP_REPLY;
                    arp_reply->id = ARP_ID;
                    memcpy(arp_reply->sender_mac, my_mac_addr, IF_HADDR);
                    memcpy(arp_reply->sender_ip, my_ip_addr, IP_SIZE);
                    memcpy(arp_reply->target_mac, arp_msg->sender_mac, IF_HADDR);
                    memcpy(arp_reply->target_ip, arp_msg->sender_ip, IP_SIZE);

                    /* Populate the ethernet header */
                    memcpy(eh_send->h_dest, eh_recv->h_source, ETH_ALEN);
                    memcpy(eh_send->h_source, my_mac_addr, ETH_ALEN);
                    eh_send->h_proto = htons(ETH_ARP_PROTO);

                    /* Send the ARP reply */
                    if(sendto(packet_sockfd, send_buf, ETH_FRAME_LEN, 0, (struct sockaddr *)&unicast_pfaddr, sizeof(unicast_pfaddr)) < 0) {
                            fprintf(stderr, "ARP: Error in sending ARP reply - %s\n", strerror(errno));
                            goto out;
                    }
                    printf("\nSending ARP_REPLY message...\n");
                    printf("\nARP_REPLY message sent from %s(%s)::\n", vm_name, inet_ntop(AF_INET, my_ip_addr, src_ipaddress, INET_ADDRSTRLEN));
                    printf("Ethernet frame header -- src: %02x:%02x:%02x:%02x:%02x:%02x --> dst: %02x:%02x:%02x:%02x:%02x:%02x\n", 
                           (unsigned char)eh_send->h_source[0], (unsigned char)eh_send->h_source[1], (unsigned char)eh_send->h_source[2], (unsigned char)eh_send->h_source[3], (unsigned char)eh_send->h_source[4], 
                           (unsigned char)eh_send->h_source[5],
                           (unsigned char)eh_send->h_dest[0], (unsigned char)eh_send->h_dest[1], (unsigned char)eh_send->h_dest[2], (unsigned char)eh_send->h_dest[3], (unsigned char)eh_send->h_dest[4],
                           (unsigned char)eh_send->h_dest[5]);

                    printf("Payload -- sender_ip: %s, sender_mac: %02x:%02x:%02x:%02x:%02x:%02x, target_ip: %s, target_mac: %02x:%02x:%02x:%02x:%02x:%02x\n\n", 
                           inet_ntop(AF_INET, arp_reply->sender_ip, src_ipaddress, INET_ADDRSTRLEN), 
                           (unsigned char)arp_reply->sender_mac[0], (unsigned char)arp_reply->sender_mac[1], (unsigned char)arp_reply->sender_mac[2], (unsigned char)arp_reply->sender_mac[3], (unsigned char)arp_reply->sender_mac[4], (unsigned char)arp_reply->sender_mac[5],
                           inet_ntop(AF_INET, arp_reply->target_ip, dest_ipaddress, INET_ADDRSTRLEN),
                           (unsigned char)arp_reply->target_mac[0], (unsigned char)arp_reply->target_mac[1], (unsigned char)arp_reply->target_mac[2], (unsigned char)arp_reply->target_mac[3], (unsigned char)arp_reply->target_mac[4], (unsigned char)arp_reply->target_mac[5]);

                    /* Insert entry into cache */
                    if((temp = already_in_cache(arp_msg->sender_ip))) {

                            /* Entry already exists; simply update the values */
                            memcpy(temp->if_haddr, arp_msg->sender_mac, IF_HADDR);
                            temp->if_index = if_index;
                            temp->sll_hatype = hatype;
                            // temp->domain_sockfd = -1; /* TODO: Check this */
                    } else {
                            memcpy(node->ip_addr, arp_msg->sender_ip, IP_SIZE);
                            memcpy(node->if_haddr, arp_msg->sender_mac, IF_HADDR);
                            node->if_index = if_index;
                            node->sll_hatype = hatype;
                            node->domain_sockfd = -1;

                            list_add_tail(&node->list_element, &cache_list_head);
                    }

            }
            else {
                    // printf("\nDEBUG: I am not the destination node.\n");

                    if((temp = already_in_cache(arp_msg->sender_ip))) {

                            /* Entry already exists; simply update the values */
                            memcpy(temp->if_haddr, arp_msg->sender_mac, IF_HADDR);
                            temp->if_index = if_index;
                            temp->sll_hatype = hatype;
                    } 
            }
    }
    else if(arp_msg->type == ARP_REPLY) {
            if((temp = already_in_cache(arp_msg->sender_ip))) {
                    struct hwaddr haddr;
                    char src_ipaddress[INET_ADDRSTRLEN + 1] = {0};
                    char dest_ipaddress[INET_ADDRSTRLEN + 1] = {0};
                    char vm_name[128] = {0};

                    gethostname(vm_name, 128);
                    printf("\nReceived ARP_REPLY message...\n");
                    printf("\nARP_REPLY message received at %s(%s)::\n", vm_name, inet_ntop(AF_INET, my_ip_addr, src_ipaddress, INET_ADDRSTRLEN));
                    printf("Ethernet frame header -- src: %02x:%02x:%02x:%02x:%02x:%02x --> dst: %02x:%02x:%02x:%02x:%02x:%02x\n", 
                           (unsigned char)eh_recv->h_source[0], (unsigned char)eh_recv->h_source[1], (unsigned char)eh_recv->h_source[2], (unsigned char)eh_recv->h_source[3], (unsigned char)eh_recv->h_source[4], 
                           (unsigned char)eh_recv->h_source[5],
                           (unsigned char)eh_recv->h_dest[0], (unsigned char)eh_recv->h_dest[1], (unsigned char)eh_recv->h_dest[2], (unsigned char)eh_recv->h_dest[3], (unsigned char)eh_recv->h_dest[4],
                           (unsigned char)eh_recv->h_dest[5]);

                    printf("Payload -- sender_ip: %s, sender_mac: %02x:%02x:%02x:%02x:%02x:%02x, target_ip: %s, target_mac: %02x:%02x:%02x:%02x:%02x:%02x\n\n", 
                           inet_ntop(AF_INET, arp_msg->sender_ip, src_ipaddress, INET_ADDRSTRLEN), 
                           (unsigned char)arp_msg->sender_mac[0], (unsigned char)arp_msg->sender_mac[1], (unsigned char)arp_msg->sender_mac[2], (unsigned char)arp_msg->sender_mac[3], (unsigned char)arp_msg->sender_mac[4], (unsigned char)arp_msg->sender_mac[5],
                           inet_ntop(AF_INET, arp_msg->target_ip, dest_ipaddress, INET_ADDRSTRLEN),
                           (unsigned char)arp_msg->target_mac[0], (unsigned char)arp_msg->target_mac[1], (unsigned char)arp_msg->target_mac[2], (unsigned char)arp_msg->target_mac[3], (unsigned char)arp_msg->target_mac[4], (unsigned char)arp_msg->target_mac[5]);

                    /* Entry already exists; simply update the values */
                    memcpy(temp->if_haddr, arp_msg->sender_mac, IF_HADDR);
                    temp->if_index = if_index;
                    temp->sll_hatype = hatype;

                    haddr.sll_ifindex = temp->if_index;
                    haddr.sll_hatype = temp->sll_hatype;
                    haddr.sll_halen = IF_HADDR;
                    memcpy(haddr.sll_addr, temp->if_haddr, IF_HADDR);

                    if(write(temp->domain_sockfd, &haddr, sizeof(haddr)) <= 0) {
                            struct address_cache *node, *tmp;

                            // printf("\nDEBUG: Writing to areq() failed; deleting entry from cache\n");
                            list_for_each_entry_safe(node, tmp, &cache_list_head, list_element) {
                                    if(!memcmp(node->ip_addr, temp->ip_addr, IP_SIZE)) {
                                            list_del(&node->list_element);
                                            free(node);
                                    }
                            }
                    }
                    close(temp->domain_sockfd);
            } 
    }

out:
    free(recv_buf);
    free(send_buf);

}

void handle_message_on_domain_socket(int dom_sockfd, int packet_sockfd, struct sockaddr_ll *pfaddr)
{
        int connfd;
        struct sockaddr_un cliaddr;
        socklen_t len = sizeof(cliaddr);
        int nbytes = 0;
        struct sockaddr_in *cli_info;
        unsigned char * send_buf = malloc(ETH_FRAME_LEN);
        struct ethhdr *eh_send;
        struct arp_packet *arp_request;
        struct address_cache *node;
        char ip_addr[IP_SIZE] = {0};
        char *recv_buf;
        struct hwaddr *HWaddr;

        if(!send_buf) {
                fprintf(stderr, "ARP: Error allocating memory for Ethernet frame\n");
                goto out;
        }

        eh_send = (struct ethhdr *)send_buf;

        if((connfd = accept(dom_sockfd, (struct sockaddr *)&cliaddr, &len)) < 0) {
                fprintf(stderr, "ARP: Error in accepting a connection - %s\n", strerror(errno));
                goto out;
        }

        recv_buf = malloc(sizeof(struct hwaddr) + sizeof(struct sockaddr));

        if((nbytes = read(connfd, recv_buf, sizeof(struct hwaddr) + sizeof(struct sockaddr))) < 0) {
                fprintf(stderr, "ARP: Error reading from UNIX domain socket - %s\n", strerror(errno));
                goto out;
        } 
        
        // printf("\nDEBUG: Data received on domain socket\n");

        cli_info = (struct sockaddr_in *)recv_buf;
        HWaddr = (struct hwaddr *)(recv_buf + sizeof(struct sockaddr_in));

        /* Check if the required HW address is already in the cache */

        memcpy(ip_addr, &cli_info->sin_addr.s_addr, IP_SIZE);

        node = already_in_cache(ip_addr);

        if(node) {
                if(memcmp(node->if_haddr, "\0\0\0\0\0\0", ETH_ALEN)){

                        /* HW address is already available */
                        printf("\nHW address already available in cache; sending ARP_REPLY...\n");
                        HWaddr->sll_ifindex = node->if_index;
                        HWaddr->sll_hatype = node->sll_hatype;
                        HWaddr->sll_halen = IF_HADDR;
                        memcpy(HWaddr->sll_addr, node->if_haddr, IF_HADDR);
                        node->domain_sockfd = connfd;
                        if(write(node->domain_sockfd, HWaddr, sizeof(struct hwaddr)) < 0) {
                                fprintf(stderr, "ARP: Error writing to socket %d - %s\n", node->domain_sockfd, strerror(errno));
                                goto out;
                        }
                        close(node->domain_sockfd);
                } else {
                        close(connfd);
                        // node->domain_sockfd = connfd;
                }
        } else {
                char vm_name[128] = {0};
                char src_ipaddress[INET_ADDRSTRLEN + 1] = {0};
                char dest_ipaddress[INET_ADDRSTRLEN + 1] = {0};
                struct address_cache *new_node = malloc(sizeof(struct address_cache));

                memcpy(new_node->ip_addr, ip_addr, IP_SIZE);
                memset(new_node->if_haddr, 0, IF_HADDR);
                new_node->if_index = HWaddr->sll_ifindex;
                new_node->sll_hatype = HWaddr->sll_hatype;
                new_node->domain_sockfd = connfd;

                list_add_tail(&new_node->list_element, &cache_list_head);

                arp_request = (struct arp_packet *)(send_buf + 14);

                arp_request->type = ARP_REQUEST;
                arp_request->id = ARP_ID;

                memcpy(arp_request->sender_ip, my_ip_addr, IP_SIZE);
                memcpy(arp_request->sender_mac, my_mac_addr, IF_HADDR);
                memcpy(arp_request->target_ip, ip_addr, IP_SIZE);
                memset(arp_request->target_mac, 0, IF_HADDR);

                memset(eh_send->h_dest, 0xff, ETH_ALEN);
                memcpy(eh_send->h_source, my_mac_addr, ETH_ALEN);
                eh_send->h_proto = htons(ETH_ARP_PROTO);

                gethostname(vm_name, 128);
                printf("\nSending ARP_REQUEST message...\n");
                printf("\nARP_REQUEST message sent from %s(%s)::\n", vm_name, inet_ntop(AF_INET, my_ip_addr, src_ipaddress, INET_ADDRSTRLEN));
                printf("Ethernet frame header -- src: %02x:%02x:%02x:%02x:%02x:%02x --> dst: %02x:%02x:%02x:%02x:%02x:%02x\n", 
                       (unsigned char)eh_send->h_source[0], (unsigned char)eh_send->h_source[1], (unsigned char)eh_send->h_source[2], (unsigned char)eh_send->h_source[3], (unsigned char)eh_send->h_source[4], 
                       (unsigned char)eh_send->h_source[5],
                       (unsigned char)eh_send->h_dest[0], (unsigned char)eh_send->h_dest[1], (unsigned char)eh_send->h_dest[2], (unsigned char)eh_send->h_dest[3], (unsigned char)eh_send->h_dest[4],
                       (unsigned char)eh_send->h_dest[5]);

                printf("Payload -- sender_ip: %s, sender_mac: %02x:%02x:%02x:%02x:%02x:%02x, target_ip: %s, target_mac: %02x:%02x:%02x:%02x:%02x:%02x\n\n", 
                       inet_ntop(AF_INET, arp_request->sender_ip, src_ipaddress, INET_ADDRSTRLEN), 
                       (unsigned char)arp_request->sender_mac[0], (unsigned char)arp_request->sender_mac[1], (unsigned char)arp_request->sender_mac[2], (unsigned char)arp_request->sender_mac[3], (unsigned char)arp_request->sender_mac[4], (unsigned char)arp_request->sender_mac[5],
                       inet_ntop(AF_INET, arp_request->target_ip, dest_ipaddress, INET_ADDRSTRLEN),
                       (unsigned char)arp_request->target_mac[0], (unsigned char)arp_request->target_mac[1], (unsigned char)arp_request->target_mac[2], (unsigned char)arp_request->target_mac[3], (unsigned char)arp_request->target_mac[4], (unsigned char)arp_request->target_mac[5]);

                if(sendto(packet_sockfd, send_buf, ETH_FRAME_LEN, 0, (struct sockaddr *)pfaddr, sizeof(*pfaddr)) < 0) {
                        fprintf(stderr, "ARP: Error sending ARP request - %s\n", strerror(errno));
                        return;
                }

        }

out:
        free(send_buf);
}

struct address_cache * already_in_cache(char *ip_addr)
{
        struct list_head *pos;
        struct address_cache *node;

        list_for_each(pos, &cache_list_head) {
                node = list_entry(pos, struct address_cache, list_element);

                if(!memcmp(ip_addr, node->ip_addr, IP_SIZE)) {
                        /* Found entry in the cache */
                        return node;
                }
        }

        return NULL;
}

